##############################################################################
#
# test on three nodes:
#
# torchrun --nproc_per_node=1 --nnodes=3 --node_rank=0 --master_addr="172.17.0.5" --master_port=1234 /workspace/test-dist.py
# torchrun --nproc_per_node=1 --nnodes=3 --node_rank=1 --master_addr="172.17.0.5" --master_port=1234 /workspace/test-dist.py
# torchrun --nproc_per_node=1 --nnodes=3 --node_rank=2 --master_addr="172.17.0.5" --master_port=1234 /workspace/test-dist.py
#
import torch
import torch.distributed as dist
from transformers import AutoTokenizer, AutoModelForCausalLM
import logging
import os
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import Dataset, DataLoader
from torch.nn.parallel import DistributedDataParallel as DDP

# Suppress warnings
logging.getLogger("transformers").setLevel(logging.ERROR)

def use_gpu():
    if os.environ.get('NO_GPU') is not None:
        return False
    return torch.cuda.is_available()

class ProblemDataset(Dataset):
    def __init__(self, prompts):
        self.prompts = prompts

    def __len__(self):
        return len(self.prompts)

    def __getitem__(self, idx):
        return self.prompts[idx]

def setup_distributed():
    # Initialize distributed environment
    dist.init_process_group(backend='nccl' if use_gpu() else 'gloo')
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    
    # Set the device
    if use_gpu():
        if torch.cuda.get_device_capability() >= (7, 5):
            # Get the list of available GPUs
            available_gpus = list(range(torch.cuda.device_count()))
            print(f"\n************************* available_gpus:{available_gpus}")
            if not available_gpus:
                raise RuntimeError("No GPU devices available")
            device = torch.device(f"cuda:{available_gpus[rank % len(available_gpus)]}")
            torch.cuda.set_device(device)
        else:
            raise RuntimeError("No GPU devices available.")
    else:
        device = torch.device("cpu")

    return rank, world_size, device

def cleanup_distributed():
    dist.destroy_process_group()

def main():
    # Setup distributed environment
    rank, world_size, device = setup_distributed()

    # 加载模型和分词器
    model_path = "/workspace/Qwen2-0.5B/"
    print(f"Loading model from: {model_path}")
 
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.float16 if use_gpu() else torch.float32,
##        device_map="auto"  # Automatically maps model to available devices
    )

    model = model.to(device)  # Move model to device
##    model.parallelize()  # Manually parallelize the model
    # 根据设备类型决定是否指定 device_ids
    model = DDP(model, device_ids=[rank]) if use_gpu() else DDP(model)

    # 创建包含100个问题的数据集
    prompts = [
    "问题1: 请用一句话回答，中国的首都是哪里？",
    "问题2: 请用一句话回答，《红楼梦》的作者是谁？",
    "问题3: 请用一句话回答，太阳系中离太阳最近的行星是哪个？",
    "问题4: 请用一句话回答，人体最大的器官是什么？",
    "问题5: 请用一句话回答，1公里等于多少米？",
    "问题6: 请用一句话回答，世界上面积最大的国家是哪个？",
    "问题7: 请用一句话回答，水的化学式是什么？",
    "问题8: 请用一句话回答，美国独立宣言是在哪一年签署的？",
    "问题9: 请用一句话回答，计算机科学中，CPU的全称是什么？",
    "问题10: 请用一句话回答，植物进行光合作用的主要器官是什么？",
    "问题11: 请用一句话回答，中国古代四大发明是什么？",
    "问题12: 请用一句话回答，地球上最长的河流是哪条？",
    "问题13: 请用一句话回答，《蒙娜丽莎》的作者是谁？",
    "问题14: 请用一句话回答，电力的单位是什么？",
    "问题15: 请用一句话回答，联合国总部设在哪个城市？",
    "问题16: 请用一句话回答，人体有多少对染色体？",
    "问题17: 请用一句话回答，数学中，圆的周长公式是什么？",
    "问题18: 请用一句话回答，世界上使用人数最多的语言是什么？",
    "问题19: 请用一句话回答，第一次世界大战爆发于哪一年？",
    "问题20: 请用一句话回答，互联网上常用的HTTP协议的全称是什么？",
    "问题21: 请用一句话回答，中国的国旗上有几颗星星？",
    "问题22: 请用一句话回答，地球上最大的海洋是哪个？",
    "问题23: 请用一句话回答，莎士比亚的四大悲剧是哪四部？",
    "问题24: 请用一句话回答，牛顿第一运动定律又称为什么定律？",
    "问题25: 请用一句话回答，世界上最高的山峰是哪座？",
    "问题26: 请用一句话回答，中国的国宝动物是什么？",
    "问题27: 请用一句话回答，计算机存储容量的基本单位是什么？",
    "问题28: 请用一句话回答，光合作用的产物主要是什么？",
    "问题29: 请用一句话回答，美国现任总统是谁？",
    "问题30: 请用一句话回答，一年有多少个节气？",
    "问题31: 请用一句话回答，《论语》记录的是哪位思想家的言行？",
    "问题32: 请用一句话回答，地球上最深的海洋是哪个？",
    "问题33: 请用一句话回答，电灯的发明者是谁？",
    "问题34: 请用一句话回答，人体最坚硬的部位是什么？",
    "问题35: 请用一句话回答，数学中，勾股定理的公式是什么？",
    "问题36: 请用一句话回答，世界上人口最多的国家是哪个？",
    "问题37: 请用一句话回答，中国的母亲河是哪条河？",
    "问题38: 请用一句话回答，计算机编程中，'Hello World'通常用来做什么？",
    "问题39: 请用一句话回答，植物通过什么部位吸收水分？",
    "问题40: 请用一句话回答，第二次世界大战结束于哪一年？",
    "问题41: 请用一句话回答，中国的国花是什么？",
    "问题42: 请用一句话回答，太阳系中最大的行星是哪个？",
    "问题43: 请用一句话回答，《史记》的作者是谁？",
    "问题44: 请用一句话回答，力的国际单位是什么？",
    "问题45: 请用一句话回答，欧盟使用的统一货币是什么？",
    "问题46: 请用一句话回答，人体血液的主要功能是什么？",
    "问题47: 请用一句话回答，1公顷等于多少平方米？",
    "问题48: 请用一句话回答，世界上最小的国家是哪个？",
    "问题49: 请用一句话回答，互联网的发明者是谁？",
    "问题50: 请用一句话回答，绿色植物进行光合作用需要什么气体？",
    "问题51: 请用一句话回答，中国历史上第一个皇帝是谁？",
    "问题52: 请用一句话回答，地球上最大的洲是哪个？",
    "问题53: 请用一句话回答，《战争与和平》的作者是谁？",
    "问题54: 请用一句话回答，电流的单位是什么？",
    "问题55: 请用一句话回答，国际奥林匹克委员会总部设在哪个城市？",
    "问题56: 请用一句话回答，人体最大的骨骼是哪块？",
    "问题57: 请用一句话回答，数学中，π的值大约是多少？",
    "问题58: 请用一句话回答，世界上使用最广泛的字母系统是什么？",
    "问题59: 请用一句话回答，美国宪法是哪一年颁布的？",
    "问题60: 请用一句话回答，计算机科学中，RAM的全称是什么？",
    "问题61: 请用一句话回答，中国的国歌叫什么名字？",
    "问题62: 请用一句话回答，地球上最热的沙漠是哪个？",
    "问题63: 请用一句话回答，《哈姆雷特》的作者是谁？",
    "问题64: 请用一句话回答，重力的发现者是谁？",
    "问题65: 请用一句话回答，世界上最大的岛屿是哪个？",
    "问题66: 请用一句话回答，中国的四大名著是什么？",
    "问题67: 请用一句话回答，计算机存储中，1GB等于多少MB？",
    "问题68: 请用一句话回答，植物生长需要的主要营养元素有哪些？",
    "问题69: 请用一句话回答，现任联合国秘书长是谁？",
    "问题70: 请用一句话回答，中国传统历法是什么历？",
    "问题71: 请用一句话回答，儒家思想的创始人是谁？",
    "问题72: 请用一句话回答，地球上最长的山脉是哪个？",
    "问题73: 请用一句话回答，电话的发明者是谁？",
    "问题74: 请用一句话回答，人体最重要的解毒器官是什么？",
    "问题75: 请用一句话回答，数学中，一元二次方程的求根公式是什么？",
    "问题76: 请用一句话回答，世界上最大的湖泊是哪个？",
    "问题77: 请用一句话回答，中国的国球是什么运动？",
    "问题78: 请用一句话回答，计算机编程中，面向对象编程的三大特性是什么？",
    "问题79: 请用一句话回答，植物通过什么过程释放氧气？",
    "问题80: 请用一句话回答，冷战结束的标志性事件是什么？",
    "问题81: 请用一句话回答，中国最大的淡水湖是哪个？",
    "问题82: 请用一句话回答，太阳系中唯一有生命存在的行星是哪个？",
    "问题83: 请用一句话回答，《三国演义》的作者是谁？",
    "问题84: 请用一句话回答，电压的单位是什么？",
    "问题85: 请用一句话回答，世界贸易组织(WTO)总部设在哪个城市？",
    "问题86: 请用一句话回答，人体最小的骨骼是哪块？",
    "问题87: 请用一句话回答，数学中，黄金分割的比例大约是多少？",
    "问题88: 请用一句话回答，世界上使用最广泛的语言是什么？",
    "问题89: 请用一句话回答，工业革命最早开始于哪个国家？",
    "问题90: 请用一句话回答，计算机科学中，HTML的全称是什么？",
    "问题91: 请用一句话回答，中国有多少个省级行政区？",
    "问题92: 请用一句话回答，地球上最冷的大陆是哪个？",
    "问题93: 请用一句话回答，《西游记》的作者是谁？",
    "问题94: 请用一句话回答，能量守恒定律是谁提出的？",
    "问题95: 请用一句话回答，世界上最大的半岛是哪个？",
    "问题96: 请用一句话回答，中国的传统医学称为什么？",
    "问题97: 请用一句话回答，计算机中，二进制使用的基数是多少？",
    "问题98: 请用一句话回答，植物分类的基本单位是什么？",
    "问题99: 请用一句话回答，第一次登月成功是在哪一年？",
    "问题100: 请用一句话回答，中国共产党的成立时间是哪一年？"
    ]
    dataset = ProblemDataset(prompts)

    # 使用DistributedSampler划分数据集
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=False)
    dataloader = DataLoader(dataset, batch_size=1, sampler=sampler)

    # 处理数据集
#    count=0
    for batch_prompt in dataloader:
        prompt = batch_prompt[0]
        inputs = tokenizer(prompt, return_tensors="pt").to(device)

        output_sequences = model.module.generate(
            **inputs,
            max_new_tokens=20,
            do_sample=True,
            temperature=0.7
        )

        # Extract just the generated part, not including input
        input_length = inputs.input_ids.shape[1]
        response = tokenizer.decode(output_sequences[0][input_length:], skip_special_tokens=True)

#        count = count + 1
#        print(f"\nOUTPUT:{count}")
        print(f"n{'='*50}\nINPUT:\n{prompt}\nOUTPUT:{response}")
        dist.barrier()  # 确保所有进程都完成当前 batch 的处理，避免输出交错


    # Cleanup
    cleanup_distributed()

if __name__ == "__main__":
    main()
